# StemBranchCalculator.py
from constants import BAGUA_ORDERED, BAGUA_TO_NATURE, GAN_SHU, HIDE_SHU, ZHI_SHU
from utils.format_lst import format_max3_to_max2
from mapping.numLst2idx import GuaConverter

STD = "num_lst 默认长度为6，每个数0-3"


class StemBranchCalculator:
    def __init__(self, num_lst: list[int], have_hide: bool = False):
        '''num_lst 默认长度为6，每个数0-3'''
        assert len(num_lst) == 6, STD
        assert all([r in [0, 1, 2, 3] for r in num_lst]), STD
        self.num_lst = num_lst
        self.have_hide = have_hide  # 是否包含伏藏卦的天干地支数
        self._hide_mask = None

    def to_gan(self, is_filter: bool = True) -> list[list[int]]:
        """
        已知 numLst 求天干数
        :param is_filter: 是否过滤未发动的爻
        :return: 天干数矩阵
        """
        prop_source = [[li[0]] * 3 + [li[-1]] * 3 for li in GAN_SHU]
        return self._calculate_property(prop_source, is_filter)

    def to_zhi(self, is_filter: bool = True) -> list[list[int]]:
        """
        已知 numLst 求地支数
        :param is_filter: 是否应用掩码,过滤未发动的爻
        :return: 地支数矩阵
        """
        return self._calculate_property(ZHI_SHU, is_filter)

    def _calculate_property(self, prop_source: dict, is_filter: bool = True) -> list[list[int]]:
        """
        通用属性计算方法，用于天干、地支等
        :param prop_source: 属性数据源，如 GAN_SHU 或 ZHI_SHU
        :param is_filter: 是否过滤未发动的爻
        :param apply_mask: 是否对变动卦应用掩码
        :return: 计算结果列表
        """
        result = []
        origin, changed = format_max3_to_max2(self.num_lst)

        # 获取内外卦自然属性索引
        outer_idx = GuaConverter(origin[:3]).get_nature_idx()
        inner_idx = GuaConverter(origin[3:]).get_nature_idx()

        # 构建主卦属性
        origin_lst = list(prop_source[outer_idx][:3] + prop_source[inner_idx][3:])
        result.append(origin_lst)

        if origin == changed:
            result.append([-1] * 6)
        else:
            idx3 = GuaConverter(changed[:3]).get_nature_idx()
            idx4 = GuaConverter(changed[3:]).get_nature_idx()
            changed_lst = list(prop_source[idx3][:3] + prop_source[idx4][3:])
            if is_filter:
                result.append([changed_lst[i] if self.num_lst[i] in [0, 3] else -1 for i in range(6)])
            else:
                result.append(changed_lst)

        # 处理伏藏卦
        if self.have_hide:
            bagua_idx = GuaConverter(self.num_lst).get_home_idx() // 8
            home_nature_gua_idx = BAGUA_TO_NATURE[bagua_idx]
            hide_lst = list(prop_source[home_nature_gua_idx])
            hide_result = [hide_lst[i] if self.hide_mask[i] else -1 for i in range(6)]
            result.append(hide_result)

        return result

    @property
    def hide_mask(self) -> list[bool]:
        '''已有缓存则返回，尚未缓存则计算'''
        if not self._hide_mask:
            self._hide_mask = self.get_hide_mask()
        return self._hide_mask

    def get_hide_mask(self) -> list[bool]:
        '''根据项目卦序，计算伏藏卦的掩码'''
        hide_mask = [False] * 6
        gua_idx, _ = GuaConverter(self.num_lst).get_gua_idx(is_pair=True)
        if gua_idx in HIDE_SHU:
            for idx in HIDE_SHU[gua_idx]:
                hide_mask[idx] = True
        return hide_mask

    def get_mask(self) -> list[list[bool]]:
        """
        获取掩码矩阵
        :return: 掩码矩阵
        """
        mask = []
        # 主卦掩码全为True
        mask.append([True] * 6)
        # 变卦的掩码看是否发动
        changed_mask = [True if self.num_lst[i] in [0, 3] else False for i in range(6)]
        mask.append(changed_mask)
        if self.have_hide: mask.append(self.hide_mask)
        return mask


if __name__ == '__main__':
    coin_num_lst = [1, 0, 2, 1, 0, 3]
    sbc = StemBranchCalculator(coin_num_lst, True)
    print(sbc.to_gan())
    print(sbc.to_zhi())
